package com.bruce.ai.config;

import com.bruce.ai.advisor.LoggingAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.Resource;

@Configuration
public class ChatClientConfig {

    /**
     * 配置ChatClient
     * @param chatClientBuilder
     * @return
     */
    @Bean
    public ChatClient getChatClient(ChatClient.Builder chatClientBuilder, ChatMemory chatMemory) {
        ChatClient chatClient = chatClientBuilder.defaultSystem("""
                        您是莱康航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                        					   您正在通过在线聊天系统与客户互动。
                        					   在提供有关预订或取消预订的信息之前，您必须始终从用户处获取以下信息：预订号、客户姓名。
                        					   在询问用户之前，请检查消息历史记录以获取此信息。
                        					   在更改预订之前，您必须确保条款允许这样做。
                        					   在执行取消机票预定之前，您必须展示退订的详细信息，并继续之前征得用户同。
                        					   如果更改需要收费，您必须在继续之前征得用户同意。
                        					   使用提供的功能获取预订详细信息、更改预订和取消预订。
                        					   如果需要，可以调用相应函数调用完成辅助动作。
                        					   请讲中文。
                        					   今天的日期是 {current_date}.
                        """)
                .defaultAdvisors(
                         new PromptChatMemoryAdvisor(chatMemory), // 添加 PromptChatMemoryAdvisor 记录聊天对话的上下文
                         new LoggingAdvisor() // 添加 LoggingAdvisor 记录聊天对话的详细过程
                ).defaultFunctions("getBookingDetails")
                .build();
        return chatClient;
    }


    /**
     * 配置向量数据库
     * @param embeddingModel
     * @param vectorStore
     * @param termsOfServiceDocs
     * @return
     */
    @Bean
    CommandLineRunner ingestTermOfServiceToVectorStore(EmbeddingModel embeddingModel, VectorStore vectorStore,
                                                       @Value("classpath:rag/terms-of-service.txt") Resource termsOfServiceDocs) {

        return args -> {
            // Ingest the document into the vector store
            vectorStore.write(                                  // 3.写入
                    new TokenTextSplitter().transform(          // 2.转换
                            new TextReader(termsOfServiceDocs).read())  // 1.读取
            );

        };
    }

    /**
     * 配置ChatMemory 记录聊天对话的上下文
     * @return
     */
    @Bean
    public ChatMemory chatMemory() {
        return new InMemoryChatMemory();
    }

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return new SimpleVectorStore(embeddingModel);
    }
}
